#include "OctreeHolder.h"

// Default constructor. Creates an empty octree.
OctreeHolder::OctreeHolder()
{
   mass = 0;
   nBodies = 0;
   length = 1.0;

   center[X] = 0.0;
   center[Y] = 0.0;
   center[Z] = 0.0;

   centerOfMass[X] = 0.0;
   centerOfMass[Y] = 0.0;
   centerOfMass[Z] = 0.0;

   body = NULL;
   params = NULL;

   for ( int i (0); i < 8; i++ ) children[i] = NULL;
}

// Constructor which creates an Octree with a given length.
// Each dimension of the new tree runs from -length to length.
OctreeHolder::OctreeHolder(double l)
{
   mass = 0;
   nBodies = 0;
   length = l;

   center[X] = 0.0;
   center[Y] = 0.0;
   center[Z] = 0.0;

   centerOfMass[X] = 0.0;
   centerOfMass[Y] = 0.0;
   centerOfMass[Z] = 0.0;

   body = NULL;
   params= NULL;

   for (int i=0; i < 8; i++)
      children[i] = NULL;
}

// Creates an octree from input datafile
OctreeHolder::OctreeHolder(string fn, ParamHolder* par)
{
   // Create an empty octree.
   mass = 0;
   nBodies = 0;
   length = 1.0;

   center[X] = 0.0;
   center[Y] = 0.0;
   center[Z] = 0.0;

   centerOfMass[X] = 0.0;
   centerOfMass[Y] = 0.0;
   centerOfMass[Z] = 0.0;

   body = NULL;

   for ( int i (0); i < 8; i++ ) children[i] = NULL;

   // Read in from file.
   params = par;
   readFromFile(fn);
}

// Creates a sub-octree from center of octree.
OctreeHolder::OctreeHolder(double cent[], double l)
{
   mass = 0;
   nBodies = 0;
   length = l;

   center[X] = cent[X];
   center[Y] = cent[Y];
   center[Z] = cent[Z];

   centerOfMass[X] = 0.0;
   centerOfMass[Y] = 0.0;
   centerOfMass[Z] = 0.0;

   body = NULL;

   for ( int i (0); i < 8; i++ ) children[i] = NULL;
}

// Destructor for OctreeHolder.
// Deletes the child octrees, which have been dynamically allocated.
// Does not delete the bodies, even though they have also been
// dynamically allocated, since it is presumed that bodies are not
// destroyed over the course of the simulation.
OctreeHolder::~OctreeHolder()
{
   for (int i = 0; i < 8; i++) {
      if (children[i] != NULL) {
         delete children[i];
	 }
      }
}

// Reads in a file containing the bodies.
// The format is "mass, pos[X], pos[y], pos[z]
void OctreeHolder::readFromFile(string fn)
{
   nBodies = 0;

   ifstream in;
   in.open(fn.c_str());

   string line;
   double nm, nx, ny, nz, nvx, nvy, nvz;

   while (in.good()) {
      getline(in, line);
      line = prepLine(line);

      if (line.length() > 0) {
         nm = getNum(line);

         line = rmFirstNum(line);

         nx = getNum(line);
         line = rmFirstNum(line);

         ny = getNum(line);
         line = rmFirstNum(line);

         nz = getNum(line);
         line = rmFirstNum(line);

         nvx = getNum(line);
         line = rmFirstNum(line);

         nvy = getNum(line);
         line = rmFirstNum(line);

         nvz = getNum(line);
         line = rmFirstNum(line);

         vec nPos = vec(3);
         vec nVel = vec(3);

         nPos[X] = nx;
         nPos[Y] = ny;
         nPos[Z] = nz;

         nVel[X] = nvx;
         nVel[Y] = nvy;
         nVel[Z] = nvz;


         Body* b = new Body(nm, nPos, nVel);

         addBody(b);
         }
    } // End loop over lines.

    centerSystem();
    rebuildTree();
}

// Writes the current state of the simulation to a text file.
// The text file is comma delimited and human readable, but can
// also be used to initialize the simulation.
//
// The format is "mass, pos[X], pos[y], pos[z]
void OctreeHolder::writeToFile(string fn)
{
   ofstream of;
   of.open(fn.c_str());
   
   Body* b = NULL;

   // The width of the output should be a parameter.
   for(int i = 0; i < nBodies; i++) {
      b = getBody(i);

      of << setw(15) << b -> mass << ",";

      of << setw(15) << b -> pos[X] << ",";
      of << setw(15) << b -> pos[Y] << ",";
      of << setw(15) << b -> pos[Z] << ",";

      of << setw(15) << b -> vel[X] << ",";
      of << setw(15) << b -> vel[Y] << ",";
      of << setw(15) << b -> vel[Z] << endl;
      }

   of.close();
}

// Add a body to the tree. If there is already a body in the
// octant then removes that body from octant, sub-divides the
// octant into eight more octants and places both particles
// in separate boxes. Continues to sub-divide octant until
// there is only one body per box
void OctreeHolder::addBody(Body* p)
{
   if (p == NULL)
      cout << "Adding a null particle." << endl;

   // If this particle lies outside the tree, resize the tree.
   for (int i = 0; i < 3; i++) {
      if (p -> pos[i] >= (center[i] + length)) {
         length = fabs(p->pos[i]) * 1.5;
         rebuildTree();
         }

      else if (p -> pos[i] <= (center[i] - length)) {
         length = fabs(p->pos[i]) * 1.5;
         rebuildTree();
         }
      }

   // There is already at least one body in this box. Add to a subbox.
   if(nBodies == 1)
   {
      OctreeHolder* dest = check_pos(p);
      dest -> addBody(p);

      dest = check_pos(body);
      dest -> addBody(body);

      body = NULL;
   }

   else if (nBodies > 1) {
      OctreeHolder* dest = check_pos(p);
      dest -> addBody(p);
   }

   // There is nothing in this box. Add it to this box.
   else {
      body = p;
      }

   centerOfMass[X] = (mass*centerOfMass[X] + p -> mass * p ->pos[X])/(mass + p -> mass);
   centerOfMass[Y] = (mass*centerOfMass[Y] + p -> mass * p ->pos[Y])/(mass + p -> mass);
   centerOfMass[Z] = (mass*centerOfMass[Z] + p -> mass * p ->pos[Z])/(mass + p -> mass);

   mass += p -> mass;
   nBodies++;
}

// Rebuild the octree.
// This function ensures that each particle is within the correct
// box, and that no particles lie outside of the simulation.
// This function must be called after each particle is added
// and after each timestep.
void OctreeHolder::rebuildTree()
{
   // Figure out if the octree needs to be resized.
   double maxDim = 0.0;

   for (int i = 0; i < nBodies; i++) {
      Body* b = getBody(i);

      if (fabs(b -> pos[X]) > maxDim)
         maxDim = fabs(b -> pos[X]);
      if (fabs(b -> pos[Y]) > maxDim)
         maxDim = fabs(b -> pos[X]);
      if (fabs(b -> pos[Z]) > maxDim)
         maxDim = fabs(b -> pos[X]);
       }

   // Create a new Octree. Make sure that it can hold all of the
   // particles.
   OctreeHolder* newTree;
   
   if (maxDim < length)
      newTree = new OctreeHolder(length);

   else
      newTree = new OctreeHolder(maxDim*1.5);

   newTree -> setParams(params);

   // Add all of the particles to the new octree.
   for (int i = 0; i < nBodies; i++) {
      Body* b = getBody(i);
      newTree -> addBody(b);
      }

   // Copy the new tree to this tree.
   mass = newTree -> getMass();
   nBodies = newTree -> size();
   length = newTree -> getLength();
   center[X] = newTree -> center[X];
   center[Y] = newTree -> center[Y];
   center[Z] = newTree -> center[Z];
   setParams(newTree -> params);

   body = NULL;
   body = newTree -> getTopBody();

   for ( int i (0); i < 8; i++ ) {
      if (newTree -> children[i] != NULL)
         children[i] = newTree -> children[i];
      else
         children[i] = NULL;
      }
}

// check_pos
OctreeHolder* OctreeHolder::check_pos( Body * p )
{
   double l(length), w(length), h(length);
   double cent[3];

   if( center[X] <= p->pos[X] ) {
   
      if( center[Y] <= p -> pos[Y] )
      {
         if( center[Z] <= p -> pos[Z] )
         {
	    if (children[0] == NULL) {
	       cent[X] = center[X] + length/2.0;
	       cent[Y] = center[Y] + length/2.0;
	       cent[Z] = center[Z] + length/2.0;

               children[0] = new OctreeHolder( cent, l/2.0);
	       children[0] -> setParams(params);
	    }
            return children[0];
         }
         else if( p->pos[Z] < center[Z] )
         {
	    if (children[4] == NULL) {
	       cent[X] = center[X] + length/2.0;
	       cent[Y] = center[Y] + length/2.0;
	       cent[Z] = center[Z] - length/2.0;

               children[4] = new OctreeHolder(cent, l/2.0);
	       children[4] -> setParams(params);
	       }
            return children[4];
         }
      }
      else if( p -> pos[Y] < center[Y] )
      {
         if( center[Z] <= p -> pos[Z] )
         {  
	    if (children[1] == NULL) {
	       cent[X] = center[X] + length/2.0;
	       cent[Y] = center[Y] - length/2.0;
	       cent[Z] = center[Z] + length/2.0;

               children[1] = new OctreeHolder( cent, l/2.0);
	       children[1] -> setParams(params);
	       }

            return children[1];
         }
         else if( p -> pos[Z] < center[Z] )
         {
	    if (children[5] == NULL) {
	       cent[X] = center[X] + length/2.0;
	       cent[Y] = center[Y] - length/2.0;
	       cent[Z] = center[Z] - length/2.0;

               children[5] = new OctreeHolder( cent, l/2.0);
	       children[5] -> setParams(params);
	       }

            return children[5];
         }
      }
   }
   else if( p -> pos[X] < center[X] )
   {
      if( center[Y] <= p -> pos[Y] )
      {
         if( center[Z] <= p->pos[Z] )
         {
	    if (children[3] == NULL) {
	       cent[X] = center[X] - length/2.0;
	       cent[Y] = center[Y] + length/2.0;
	       cent[Z] = center[Z] + length/2.0;

               children[3] = new OctreeHolder( cent, l/2.0);
	       children[3] -> setParams(params);
	       }

            return children[3];
         }
         else if( p->pos[Z] < center[Z] )
         {
	    if (children[7] == NULL) {
	       cent[X] = center[X] - length/2.0;
	       cent[Y] = center[Y] + length/2.0;
	       cent[Z] = center[Z] - length/2.0;

               children[7] = new OctreeHolder(cent, l/2.0);
	       children[7] -> setParams(params);
	       }

            return children[7];
         }
      }
      else if( p->pos[Y] < center[Y] )
      {
         if( center[Z] <= p->pos[Z] )
         {
	    if (children[2] == NULL) {
	       cent[X] = center[X] - length/2.0;
	       cent[Y] = center[Y] - length/2.0;
	       cent[Z] = center[Z] + length/2.0;

               children[2] = new OctreeHolder( cent, l/2.0);
	       children[2] -> setParams(params);
	       }

            return children[2];
         }
         else if( p->pos[Z] < center[Z] )
         {
	    if (children[6] == NULL) {
	       cent[X] = center[X] - length/2.0;
	       cent[Y] = center[Y] - length/2.0;
	       cent[Z] = center[Z] - length/2.0;

               children[6] = new OctreeHolder( cent, l/2.0);
	       children[6] -> setParams(params);
	       }

            return children[6];
         }
      }
   }
}

// Treating the holder as if it were an array, returns the body at n
Body* OctreeHolder::getBody(int n)
{
   if (n >= nBodies) {
      cout << "Error! Trying to get a body that isn't in this box." << endl;
      cout << "Leaving getBody..." << endl;
      return NULL;
      }

   int counter = 0;

   // Is this the particle we're looking for?
   if (n == 0 && body!= NULL) {
      return body;
      }

   else {
      for(int i=0; i<8; i++) {
         if( children[i] != NULL )  {
            if (n >= counter && (n < counter + children[i] -> size())) {
              return children[i] -> getBody(n-counter);
	      }

            else
               counter += children[i] -> size();
            }
         }
      }

   // Should never reach here.
   cout << "Error. Should never have reached here." << endl;
   return NULL;
}

// Accessor function for the number of bodies in the simulation.
int OctreeHolder::size()
{
   return nBodies;
}

// force l/D < a then use center of box and mass of box
vec OctreeHolder::force(Body* m1)
{
   double F, dist;
   vec netForce(3);

   for(int i = 0; i < 3; i++) 
      netForce[i] = 0.0;

   double eps = params->SofteningParameter;
   double G = params->G_cgs/(pow(params->LengthUnits,3.0)
               /(params->MassUnits*params->TimeUnits
               *params->TimeUnits));
   
   if(nBodies == 1) {
     // cout << "Beep 1" << endl;
      // No force from self action.
      if (body == m1) {
         netForce[X] = 0.0;
	 netForce[Y] = 0.0;
	 netForce[Z] = 0.0;
	 }

      else {
         dist = sqrt(pow(body->pos[X] - m1->pos[X], 2.0)
            + pow(body->pos[Y] - m1->pos[Y], 2.0)
            + pow(body->pos[Z] - m1->pos[Z], 2.0));
      
         F = G * body->mass * m1->mass / (dist*dist + eps*eps);
      
         netForce[X] -= F * (m1->pos[X] - body->pos[X]) / dist;
         netForce[Y] -= F * (m1->pos[Y] - body->pos[Y]) / dist;
         netForce[Z] -= F * (m1->pos[Z] - body->pos[Z]) / dist;
         }
      }

   else if ( nBodies > 1 ) {
      dist = sqrt(pow(centerOfMass[X] - m1->pos[X], 2.0)
         + pow(centerOfMass[Y] - m1->pos[Y], 2.0)
         + pow(centerOfMass[Z] - m1->pos[Z], 2.0));

      if( (2.0*length/dist) < params -> OpeningAngle) {
   //      cout << "Beep 2" << endl;
         F = G * mass * m1->mass / (dist*dist + eps*eps);

//	 cout << "G: " << G << endl;
//	 cout << "F: " << F << endl;
//	 cout << "mass: " << mass << endl;
//	 cout << "Dist: " << dist << endl;
//	 cout << "Length: " << length << endl;
//	 cout << "COM: " << centerOfMass[X] << endl;
//	 cout << "COM y: " << centerOfMass[Y] << endl;
//	 cout << "COM z: " << centerOfMass[Z] << endl;
      
         netForce[X] -= F * ((m1->pos[X] - centerOfMass[X]) / dist);
         netForce[Y] -= F * ((m1->pos[Y] - centerOfMass[Y]) / dist);
         netForce[Z] -= F * ((m1->pos[Z] - centerOfMass[Z]) / dist);
         }

      else {
    //     cout << "Beep 3" << endl;
//	 cout << length << endl;
//	 cout << dist << endl;
  //       cout << params -> OpeningAngle << endl;
         for(int i(0); i<8; i++) {
	    if (children[i] != NULL) {
	       vec childForce(3);
               childForce = children[i] -> force(m1);

	       netForce[X] += childForce[X];
	       netForce[Y] += childForce[Y];
	       netForce[Z] += childForce[Z];
	       }
	    }
         }
      }
  
   return netForce;
}

// Compute the acceleration on body j.
vec OctreeHolder::accel(int j)
{
   vec a = vec(3);
   vec netForce = vec(3);

   Body* m1 = getBody(j);
   
   netForce = force(m1);
   
   a[X] = netForce[X]/m1->mass;
   a[Y] = netForce[Y]/m1->mass;
   a[Z] = netForce[Z]/m1->mass;

//   cout << "Accel x: " << a[X] << endl;

   return a;
}

// Removes comments and spaces
string OctreeHolder::prepLine(string line) {
   // Remove comments.
   int pos = line.find("#");

   if (pos != string::npos)
   line = line.substr(0, pos);

   pos = line.find("//");

   if (pos != string::npos)
   line = line.substr(0, pos);

   // Remove spaces
   for (int i = 0; i < line.length(); i++) {
      if (line[i] == ' ')
         line = line.substr(0, i) + line.substr(i+1, line.length());
            }

   return line;
}

// Returns the first number in a comma delimited string.
double OctreeHolder::getNum(string s)
{
   int pos = s.find_first_of(',');
   string numstr = s.substr(0, pos);

   istringstream ss(numstr);

   float num;
   ss >> num;

   return (double)num;
}

// Removes the first number in a comma delimited string. Also removes
// the comma.
string OctreeHolder::rmFirstNum(string s)
{
   int pos = s.find(',');
   return s.substr(pos + 1, s.size() - pos);
}

// Adjust the positions and velocitie so that the center of mass is
// fixed at the center of the system.
void OctreeHolder::centerSystem()
{
   vec cm = vec(3); // Compute the position of the center of mass.
   vec cv = vec(3); // Compute the velocity of the center of mass.

   double totalMass = 0.0;

   int n = nBodies;

   for (int i = 0; i < n; i++) {
      Body* b = getBody(i);

      cm[X] += b -> mass * b -> pos[X];
      cm[Y] += b -> mass * b -> pos[Y];
      cm[Z] += b -> mass * b -> pos[Z];

      cv[X] += b -> mass * b -> vel[X];
      cv[Y] += b -> mass * b -> vel[Y];
      cv[Z] += b -> mass * b -> vel[Z];

      totalMass += b -> mass;
      }

   cm[X] /= totalMass;
   cm[Y] /= totalMass;
   cm[Z] /= totalMass;

   cv[X] /= totalMass;
   cv[Y] /= totalMass;
   cv[Z] /= totalMass;

   for (int i = 0; i < n; i++) {
      Body* b = getBody(i);

      b -> pos[X] -= cm[X];
      b -> pos[Y] -= cm[Y];
      b -> pos[Z] -= cm[Z];

      b -> vel[X] -= cv[X];
      b -> vel[Y] -= cv[Y];
      b -> vel[Z] -= cv[Z];
      }
}

double OctreeHolder::getLength()
{
   return length;
}

double OctreeHolder::getMass()
{
   return mass;
}


Body* OctreeHolder::getTopBody()
{
   if (body!= NULL)
      return body;
   else
      return NULL;
}

void OctreeHolder::setParams(ParamHolder* p)
{
   if (p == NULL)
      cout << "Warning! Setting params to NULL!" << endl;

   params = p;
}
